package com.gitee.wsl.coroutines.strategy.race

import arrow.core.nonFatalOrThrow
import com.gitee.wsl.base.Or3
import com.gitee.wsl.coroutines.api.PipeAble
import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.Deferred
import kotlinx.coroutines.async
import kotlinx.coroutines.cancelAndJoin
import kotlinx.coroutines.coroutineScope
import kotlinx.coroutines.selects.select
import kotlin.contracts.ExperimentalContracts
import kotlin.contracts.InvocationKind
import kotlin.contracts.contract
import kotlin.coroutines.CoroutineContext
import kotlin.coroutines.EmptyCoroutineContext


/**
 * Races the participants [fa], [fb] & [fc] on the provided [CoroutineContext].
 * The winner of the race cancels the other participants.
 * Cancelling the operation cancels all participants.
 *
 * Coroutine context is inherited from a [CoroutineScope], additional context elements can be specified with [ctx] argument.
 * If the combined context does not have any dispatcher nor any other [ContinuationInterceptor], then [Dispatchers.Default] is used.
 * **WARNING** If the combined context has a single threaded [ContinuationInterceptor], this function will not run [fa], [fb] & [fc] in parallel.
 *
 * @see raceN for a function that ensures operations run in parallel on the [Dispatchers.Default].
 */
@OptIn(ExperimentalContracts::class)
@Suppress("LEAKED_IN_PLACE_LAMBDA")
suspend inline fun <A, B, C> PipeAble.Companion.raceN(
    ctx: CoroutineContext = EmptyCoroutineContext,
    crossinline fa: suspend CoroutineScope.() -> A,
    crossinline fb: suspend CoroutineScope.() -> B,
    crossinline fc: suspend CoroutineScope.() -> C
): Or3<A, B, C> {
    contract {
        callsInPlace(fa, InvocationKind.AT_MOST_ONCE)
        callsInPlace(fb, InvocationKind.AT_MOST_ONCE)
        callsInPlace(fc, InvocationKind.AT_MOST_ONCE)
    }
    return coroutineScope {
        val a = async(ctx) { fa() }
        val b = async(ctx) { fb() }
        val c = async(ctx) { fc() }
        select<Or3<A, B, C>> {
            a.onAwait.invoke { Or3.One(it) }
            b.onAwait.invoke { Or3.Two(it) }
            c.onAwait.invoke { Or3.Three(it) }
        }.also {
            when (it) {
                is Or3.One -> cancelAndCompose(b, c)
                is Or3.Two -> cancelAndCompose(a, c)
                is Or3.Three -> cancelAndCompose(a, b)
            }
        }
    }
}

@PublishedApi
internal suspend fun cancelAndCompose(first: Deferred<*>, second: Deferred<*>): Unit {
    val e1 = try {
        first.cancelAndJoin()
        null
    } catch (e: Throwable) {
        e.nonFatalOrThrow()
    }
    val e2 = try {
        second.cancelAndJoin()
        null
    } catch (e: Throwable) {
        e.nonFatalOrThrow()
    }
    (e1?.apply { e2?.let(::addSuppressed) } ?: e2)?.let { throw it }
}
